#include <iostream>
#include <vector>
#include <map>

#include "psUtil.h"
#include "coloring.h"

using namespace std;

Coloring::Coloring(int nInds, int nSteps, 
		   int maxColor, intMatrixType data) {
	this->nInds = nInds;
	this->nSteps = nSteps;
	this->maxColor = maxColor;
	this->data = data;

	setFactMaxColor();
}

Coloring::~Coloring() {

}

intMatrixType Coloring::getInitColoring(intVec colors) {
	intMatrixType coloring;
	for (int i = 0; i < nSteps; i++) {
		int colorIndex = 0;
		intVec col;
		map<int, int> colorPairs;
		for (int j = 0; j < nInds; j++) {
			const int dv = data[j][i];
			const int color = colors[colorIndex];

			map<int, int>::iterator itr = 
				colorPairs.find(dv);
			if(itr == colorPairs.end()) {
				col.push_back(color);
				colorPairs.insert(
					make_pair(dv, 
						  color));
				colorIndex ++;
			}
			else
				col.push_back(itr->second);
		}
		coloring.push_back(col);
	}

	return coloring;
}

intMatrixType Coloring::colorTimeStep(
	intMatrixType coloring, 
	int t, intVec colors) {

	map<int, int> colorPairs;
	int colorIndex = 0;
	for (int i = 0; i < nInds; i++) {
		int origColor = coloring[i][t];
		int color = colors[colorIndex];

		map<int, int>::iterator itr = 
			colorPairs.find(origColor);
		if(itr != colorPairs.end()) {
			coloring[i][t] = itr->second;
		}
		else {
			coloring[i][t] = color;
			colorPairs.insert(
				make_pair(origColor,
					  color));
			colorIndex ++;
		}				   
	}

	return coloring;
}


void Coloring::genAllColorings(int t, 
			       intMatrixType coloring, 
			       intVec colors) {

	if (t == nSteps - 1) {
		intMatrixType tempColoring = 
			colorTimeStep(coloring, 
				      t , colors);
		allColorings.push_back(tempColoring);
	}
	else {
		for (intMatrixType::iterator clrItr = 
			     allColorPerms.begin();
		     clrItr != allColorPerms.end();
		     clrItr++) {

			intVec nextColorSet = *clrItr;
 			coloring = colorTimeStep(coloring, 
 						 t, 
 						 colors);
			genAllColorings(t + 1, 
					coloring, 
					nextColorSet);
		}
	}
}
 
intCubeType Coloring::getAllColorings() {
	intVec colors;

	for (int i = 1; i <= maxColor; i++)
		colors.push_back(i);

	this->setAllColorPermutations(colors);

  	intMatrixType initColoring = 
 		PSUtil::transposeMat(this->getInitColoring(colors));

 	this->genAllColorings(0, initColoring, colors);

 	cout << "Number of coloring instances: " 
	     << allColorings.size() << endl;

	return allColorings;
}

intVec Coloring::permutations(int k, intVec s) {
        for (int j = 1; j < maxColor; j++) {
                k /= j;
                swap(s[(k % j) + 1], s[j - 1]);
        }

        return s;
}

void Coloring::setAllColorPermutations(intVec colors) {
	const int factN = numColorPerms;
	intVec s = colors;

	for (int k = 0; k < factN; k++) {
		s = permutations(k, s);
		allColorPerms.push_back(s);
	}
}
